import os
import sys
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS
from model.nomic_embed_text import NomicEmbedText
from model.local_deep_seek_r1 import DeepSeekR1Prompt

# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

# 加载之前保存的向量数据库
embeddings = NomicEmbedText().get_llm_model()

try:
    vector_store = FAISS.load_local(
        "rag_data/faiss_index", 
        embeddings,
        allow_dangerous_deserialization=True  # 添加此参数允许pickle反序列化
    )
    print("成功加载FAISS向量库")
    # 创建检索器
    retriever = vector_store.as_retriever(
        search_type="similarity",  # 或者使用 "mmr" 最大边际相关性搜索
        search_kwargs={"k": 5}     # 返回前3个最相关的片段
    )
except Exception as e:
    print(f"加载向量库失败: {e}")
    sys.exit(1)

# 使用DeepSeek作为大模型
llm = DeepSeekR1Prompt().get_llm_model()

# 创建更复杂的提示模板，用于SQL参数替换
prompt = ChatPromptTemplate.from_template("""你是一个SQL专家，请根据用户的查询和相关的SQL模板，生成一个完整的SQL查询。

<检索到的SQL上下文>
{context}
</检索到的SQL上下文>

用户查询: {question}

请分析用户查询的意图，从检索到的SQL中组装出一个完整的SQL查询语句，并根据用户查询中的条件替换相应的参数。例如，如果用户想查询特定网关、特定日期范围或特定状态的数据，需要替换SQL中的相应条件。

返回一个完整的、可执行的SQL查询语句。只返回SQL语句，不要有其他解释。
""")

# 创建检索-生成链
retriever_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# 示例使用方法
def query_sql(user_query):
    """
    根据用户查询，检索相关SQL并替换参数
    
    参数:
        user_query: 用户自然语言查询
        
    返回:
        生成的SQL语句
    """
    return retriever_chain.invoke(user_query)

# 测试函数
if __name__ == "__main__":
    test_query = "网关：云信留客山东移动行业 2025年2月22日到2月24日的所有成功明细"
    result = query_sql(test_query)
    print("生成的SQL:")
    print(result)
